{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ORPO微調整の例\n",
    "\n",
    "このノートブックでは、`trl`ライブラリを使用してORPO（選好確率比最適化）を実行する方法を示します。ORPOは、モデルの出力を人間の選好に合わせるためのシンプルで効果的な方法です。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 環境の設定\n",
    "\n",
    "まず、必要なライブラリをインストールし、Hugging Faceにログインします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Google Colabでの要件のインストール\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Hugging Faceへの認証\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# 便利のため、Hugging Faceのトークンを.envファイルのHF_TOKENとして環境変数に設定できます"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. モデルとデータセットのロード\n",
    "\n",
    "次に、事前学習済みモデルとデータセットをロードします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "# モデルとトークナイザーをロード\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# データセットをロード\n",
    "dataset = load_dataset(\"Anthropic/hh-rlhf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. ORPOトレーナーの設定\n",
    "\n",
    "ORPOトレーナーを設定し、トレーニングを開始します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import ORPOConfig, ORPOTrainer\n",
    "\n",
    "# ORPOの設定\n",
    "orpo_config = ORPOConfig(\n",
    "    model_name_or_path=model_name,\n",
    "    learning_rate=5e-5,\n",
    "    per_device_train_batch_size=4,\n",
    "    num_train_epochs=3,\n",
    "    logging_steps=10,\n",
    "    save_steps=100,\n",
    "    output_dir=\"./orpo_output\"\n",
    ")\n",
    "\n",
    "# ORPOトレーナーを初期化\n",
    "trainer = ORPOTrainer(\n",
    "    model=model,\n",
    "    args=orpo_config,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    eval_dataset=dataset[\"test\"],\n",
    "    tokenizer=tokenizer\n",
    ")\n",
    "\n",
    "# モデルをトレーニング\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 微調整されたモデルの保存\n",
    "\n",
    "トレーニングが完了したら、微調整されたモデルを保存します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 微調整されたモデルを保存\n",
    "trainer.save_model(\"./orpo_finetuned_model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 微調整されたモデルの評価\n",
    "\n",
    "最後に、微調整されたモデルを評価します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 評価の実行\n",
    "results = trainer.evaluate()\n",
    "print(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
